# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import glob
import logging
import os

from absl import app, flags

from coin.strategy.mm.subscription import SubscriptionRequest
from coin.feed.fastfeed.util import get_products_from_sub_req
from coin.strategy.mm.feed import get_flow_book_builder_cls
from coin.exchange.kr_rest.product.product_impl import get_exchange_product_groups


def get_available_sub_req(machine, date, flow_dir=None):
  date_str = date.strftime('%Y%m%d')
  flow_dir = flow_dir or '/remote/iosg/coin/data/flow'

  sub_reqs = [
      'Spot.Binance',
      'Spot.Binance.snapshot',
      'Spot.Binance.l1_realtime',
      'Spot.Bitfinex',
      'Spot.Bitstamp',
      'Spot.Gdax',
      'Spot.Huobi',
      'Spot.Kraken',
      'Spot.Okex.v1',
      'Spot.Okex.v3',
      'Spot.Hitbtc',
      'Spot.Bequant',
      'Spot.Korbit',
      'Spot.Gopax',
      'Spot.Quoinex.single_channel',
      'Futures.Bitmex',
      'Futures.Huobi',
      'Futures.Okex.v3',
      'Futures.Okex.v3_realtime',
      'Futures.Okex.v3_swap',
      'Futures.Kraken',
      'Futures.Deribit',
      'Futures.Binance',
      'Futures.Binance.l1_realtime',
      'Futures.Bybit',
  ]

  search_topics = set()
  avail_sub_reqs = set()
  for sub_req_str in sub_reqs:
    sub_req = SubscriptionRequest.from_str(sub_req_str)
    flow_book_builder_cls = get_flow_book_builder_cls(sub_req)
    products = get_products_from_sub_req(sub_req, date)

    for product in products:
      for topic_str in flow_book_builder_cls.get_topic_string(product, '1'):
        if topic_str in search_topics:
          continue
        search_topics.add(topic_str)

        search_str = os.path.join(flow_dir, machine, date_str, topic_str) + '*'
        files = glob.glob(search_str)
        if files:
          avail_sub_reqs.add(sub_req_str)
          break

  return list(avail_sub_reqs)


def dump_feed_cache_gen_script(version,
                               date,
                               feed_cache_dir,
                               machine_whitelist=None,
                               sub_req_whitelist=None):
  date_str = date.strftime('%Y%m%d')

  machines = [
      'feed-01.ap-northeast-2.aws',
      "feed-04.ap-northeast-1.aws",
  ]

  feed_arb_sub_req = [
      'Futures.Okex.v1',
      'Futures.Okex.v3',
      'Futures.Okex.v3_realtime',
  ]

  split_group_sub_req = [
      'Futures.Okex.v3',
      'Futures.Okex.v3_realtime',
      'Futures.Huobi',
      'Spot.Binance',
      'Spot.Binance.snapshot',
      'Spot.Binance.l1_realtime',
  ]

  crop_book_snapshot_sub_req = [
      'Futures.Huobi',
  ]

  reqs = []

  for machine in machines:
    if machine_whitelist is not None and machine not in machine_whitelist:
      continue

    sub_reqs = get_available_sub_req(machine, date)
    for sub_req in sub_reqs:
      if sub_req_whitelist is not None and sub_req not in sub_req_whitelist:
        continue

      if sub_req in feed_arb_sub_req:
        workers = ['1']
      else:
        workers = ['1', '2']

      base_req = {}
      if sub_req in crop_book_snapshot_sub_req:
        base_req['crop_book_snapshot_depth'] = 20

      if sub_req in split_group_sub_req and date.weekday() != 4:
        sub_req_obj = SubscriptionRequest.from_str(sub_req)
        for group in get_exchange_product_groups(
            sub_req_obj.market_type, sub_req_obj.exchange, sub_req_obj.api_version):
          for worker_id in workers:
            reqs.append({
                'date': date_str,
                'machine': machine,
                'sub_req': sub_req,
                'worker_id': worker_id,
                'group': group,
                **base_req,
            })
      else:
        for worker_id in workers:
          reqs.append({
              'date': date_str,
              'machine': machine,
              'sub_req': sub_req,
              'worker_id': worker_id,
              'group': 'all',
              **base_req,
          })

  if version == '1':
    for idx, req in enumerate(reqs):
      cmd_optional = []
      if 'crop_book_snapshot_depth' in req:
        cmd_optional.append('--crop_book_snapshot_depth="%d"' % req['crop_book_snapshot_depth'])
      cmd = [
          'python',
          '-m',
          'coin.feed.fastfeed.feed_cache_gen',
          '--feed_cache_dir="%s"' % feed_cache_dir,
          '--machine="%s"' % req['machine'],
          '--sub_req="%s"' % req['sub_req'],
          '--worker_id="%s"' % req['worker_id'],
          '--date="%s"' % req['date'],
          '--group="%s"' % req['group'],
          *cmd_optional,
          '--quiet'
      ]
      print('echo "%d/%d" && %s' % (idx + 1, len(reqs), ' '.join(cmd)))

  elif version == '2':
    for req in reqs:
      print(
          'bash coin/feed/fastfeed/ops/run_fastfeed.sh {date} {machine} {sub_req} {worker_id} {group}'  # noqa: E501
          .format(**req))

  elif version == '3':
    for req in reqs:
      print(
          'srun -N1 -n1 --exclusive bash coin/feed/fastfeed/ops/run_fastfeed.sh {date} {machine} {sub_req} {worker_id} {group}'  # noqa: E501
          .format(**req))

  else:
    raise Exception('Unsupported version: %s' % version)


def main(argv):
  FLAGS = flags.FLAGS
  if '-' in FLAGS.date:
    from_date_str, to_date_str = FLAGS.date.split('-')
    from_date = datetime.datetime.strptime(from_date_str, '%Y%m%d').date()
    to_date = datetime.datetime.strptime(to_date_str, '%Y%m%d').date()
  else:
    from_date = datetime.datetime.strptime(FLAGS.date, '%Y%m%d').date()
    to_date = from_date

  machine_whitelist = FLAGS.machine.split(',') if FLAGS.machine is not None else None
  sub_req_whitelist = FLAGS.sub_req.split(',') if FLAGS.sub_req is not None else None

  print('#!/bin/bash')
  cur_date = from_date
  while cur_date <= to_date:
    logging.info('Dumpung %s' % cur_date)
    dump_feed_cache_gen_script(FLAGS.version,
                               cur_date,
                               FLAGS.feed_cache_dir,
                               machine_whitelist=machine_whitelist,
                               sub_req_whitelist=sub_req_whitelist)
    cur_date += datetime.timedelta(days=1)


if __name__ == '__main__':
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)
  flags.DEFINE_string('feed_cache_dir',
                      '/remote/iosg/coin-2/buckets/feed.derived.feed_cache.coin/',
                      'Feed cache dir')
  flags.DEFINE_string('date', None, 'Date')
  flags.DEFINE_string('version', '1', '')
  flags.DEFINE_string('machine', None, '')
  flags.DEFINE_string('sub_req', None, '')
  app.run(main)
